"""
多目标围捕示例
"""

import numpy as np
import sys
import os
import argparse

sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from usv_boids.core.simulation import Environment, SimulationEngine
from usv_boids.core.entities import USV, Target
from usv_boids.utils.visualization import Visualizer, Plotter


def create_patrol_points(env_width, env_height, count=8, radius=None):
    center_x, center_y = env_width / 2, env_height / 2
    if radius is None:
        radius = min(env_width, env_height) / 2.5
        
    p_pts = []
    for i in range(count):
        angle = 2 * np.pi * i / count
        x = center_x + radius * np.cos(angle)
        y = center_y + radius * np.sin(angle)
        p_pts.append(np.array([x, y]))
        
    return p_pts


def create_formation_positions(center, count, formation_type='square', formation_radius=40.0, formation_angle_offset=0.0):
    positions = []
    center = np.array(center)
    
    if formation_type == 'square':
        # 方阵分布
        side_length = int(np.ceil(np.sqrt(count)))
        row_count = side_length
        col_count = int(np.ceil(count / row_count))
        
        spacing = formation_radius * 1.8 / max(row_count, col_count)
        
        start_x = center[0] - (col_count - 1) * spacing / 2
        start_y = center[1] - (row_count - 1) * spacing / 2
        
        cos_theta = np.cos(formation_angle_offset)
        sin_theta = np.sin(formation_angle_offset)
        rotation_matrix = np.array([[cos_theta, -sin_theta], [sin_theta, cos_theta]])
        
        count_so_far = 0
        for row in range(row_count):
            for col in range(col_count):
                if count_so_far < count:
                    pos_x = start_x + col * spacing
                    pos_y = start_y + row * spacing
                    
                    rel_x = pos_x - center[0]
                    rel_y = pos_y - center[1]
                    
                    rotated = np.dot(rotation_matrix, np.array([rel_x, rel_y]))
                    final_pos = center + rotated
                    
                    positions.append(final_pos)
                    count_so_far += 1
                else:
                    break
    
    elif formation_type == 'circle':
        for i in range(count):
            angle = formation_angle_offset + 2 * np.pi * i / count
            x = center[0] + formation_radius * np.cos(angle)
            y = center[1] + formation_radius * np.sin(angle)
            positions.append(np.array([x, y]))
            
    elif formation_type == 'arc':
        arc_angle = np.pi
        for i in range(count):
            angle = formation_angle_offset + arc_angle * i / (count - 1) if count > 1 else 0
            x = center[0] + formation_radius * np.cos(angle)
            y = center[1] + formation_radius * np.sin(angle)
            positions.append(np.array([x, y]))
            
    elif formation_type == 'line':
        line_length = formation_radius * 2
        if count > 1:
            step = line_length / (count - 1)
        else:
            step = 0
            
        start_x = center[0] - (line_length / 2) * np.cos(formation_angle_offset)
        start_y = center[1] - (line_length / 2) * np.sin(formation_angle_offset)
        
        for i in range(count):
            x = start_x + step * i * np.cos(formation_angle_offset)
            y = start_y + step * i * np.sin(formation_angle_offset)
            positions.append(np.array([x, y]))
    
    elif formation_type == 'pincer':
        side_count = int(count * 0.6)
        middle_count = count - side_count
        side_count_per_side = side_count // 2
        
        # 左翼
        for i in range(side_count_per_side):
            angle_range = np.pi / 3
            angle = formation_angle_offset - np.pi/2 + angle_range * i / (side_count_per_side - 1 if side_count_per_side > 1 else 1)
            x = center[0] + formation_radius * np.cos(angle)
            y = center[1] + formation_radius * np.sin(angle)
            positions.append(np.array([x, y]))
            
        # 中央
        for i in range(middle_count):
            angle_range = np.pi / 3
            angle = formation_angle_offset - np.pi/6 + angle_range * i / (middle_count - 1 if middle_count > 1 else 1)
            x = center[0] + formation_radius * np.cos(angle)
            y = center[1] + formation_radius * np.sin(angle)
            positions.append(np.array([x, y]))
            
        # 右翼
        for i in range(count - len(positions)):
            angle_range = np.pi / 3
            angle = formation_angle_offset + np.pi/6 + angle_range * i / (side_count - side_count_per_side - 1 if side_count - side_count_per_side > 1 else 1)
            x = center[0] + formation_radius * np.cos(angle)
            y = center[1] + formation_radius * np.sin(angle)
            positions.append(np.array([x, y]))
    
    else:
        return create_formation_positions(center, count, 'square', formation_radius, formation_angle_offset)
        
    return positions


def main():
    parser = argparse.ArgumentParser(description='多目标围捕仿真')
    parser.add_argument('--usv_count', type=int, default=24, help='无人艇数量')
    parser.add_argument('--target_count', type=int, default=3, help='目标数量')
    parser.add_argument('--fleeing', action='store_true', help='目标是否具有逃逸行为')
    parser.add_argument('--disable_capture', action='store_true', default=True, help='禁用目标捕获机制')
    parser.add_argument('--save_animation', type=str, help='保存动画到文件')
    parser.add_argument('--save_plot', type=str, help='保存指标图表到文件')
    args = parser.parse_args()
    
    print(f"启动多目标围捕 (无人艇: {args.usv_count}, 目标: {args.target_count})...")
    
    # 创建仿真环境
    env = Environment(width=150.0, height=150.0, boundary_mode='wrap')
    sim = SimulationEngine(env)
    
    # 创建巡逻点
    patrol = create_patrol_points(env.width, env.height, count=8)
    
    # 添加多个目标
    for i in range(args.target_count):
        point_idx = i % len(patrol)
        pos = patrol[point_idx].copy()
        
        next_idx = (point_idx + 1) % len(patrol)
        direction = patrol[next_idx] - pos
        
        speed = np.random.uniform(0.4, 0.7)
        vel = direction / np.linalg.norm(direction) * speed
        
        size = np.random.uniform(1.2, 2.0)
        max_spd = speed * 2
        
        target = Target(
            position=pos,
            velocity=vel,
            max_speed=max_spd,
            size=size,
            is_fleeing=args.fleeing
        )
        target.disable_capture = args.disable_capture
        sim.add_target(target)
    
    # 初始化无人艇
    if args.usv_count > 0:
        formation_center = np.array([env.width / 2, env.height / 2])
        formation_radius = min(env.width, env.height) / 4
        formation_angle = 0.0
        
        # 创建无人艇位置
        positions = create_formation_positions(
            formation_center, 
            args.usv_count, 
            formation_type='square',
            formation_radius=formation_radius,
            formation_angle_offset=formation_angle
        )
        
        # 添加所有无人艇
        for i, pos in enumerate(positions):
            angle = np.random.uniform(0, 2 * np.pi)
            speed = np.random.uniform(0.2, 0.5)
            vel = np.array([np.cos(angle), np.sin(angle)]) * speed
            
            # 创建USV
            usv = USV(
                position=pos,
                velocity=vel,
                max_speed=np.random.uniform(1.8, 2.4),
                perception_radius=np.random.uniform(30.0, 38.0),
                max_force=np.random.uniform(0.35, 0.7)
            )
            sim.add_usv(usv)
    else:
        sim.add_usv_fleet(
            count=args.usv_count,
            spawn_area=(10.0, 10.0, env.width - 10.0, env.height - 10.0),
            speed_range=(1.0, 1.8),
            perception_radius=32.0
        )
    
    # 配置无人艇
    for i, usv in enumerate(sim.usvs):
        usv.separation_weight = np.random.uniform(1.8, 2.2)
        usv.alignment_weight=np.random.uniform(1.0, 1.4)
        usv.cohesion_weight =np.random.uniform(0.8, 1.2)
        usv.target_weight = np.random.uniform(3.5, 4.5)
        
        usv.orbit_distance=np.random.uniform(3.0, 4.0)
        usv.target_switch_threshold = 0.7
        usv.steering_smoothing = 0.3
        
        usv.min_speed = 0.3
        usv.speed_adjustment_factor= 0.8
        
        if i % 3 == 0:
            # 追踪型
            usv.target_weight += 1.5
            usv.max_speed+=0.5
            usv.perception_radius=40.0
            usv.max_force += 0.4
            usv.min_speed=0.5
            usv.target_switch_threshold = 0.8
            usv.steering_smoothing=0.2
        elif i % 3 == 1:
            # 包围型
            usv.target_weight += 0.8
            usv.separation_weight +=0.5
            usv.perception_radius = 35.0
            usv.orbit_distance +=1.0
            usv.min_speed = 0.2
            usv.speed_adjustment_factor=0.6
            usv.target_switch_threshold=0.6
            usv.steering_smoothing = 0.4
        else:
            # 支援型
            usv.alignment_weight+=0.3
            usv.cohesion_weight+=0.3
            usv.perception_radius=45.0
            usv.max_speed+=0.3
            usv.min_speed=0.4  
            usv.speed_adjustment_factor = 0.7
            usv.target_switch_threshold=0.7
            usv.steering_smoothing=0.3
    
    print("执行初始目标分配...")
    sim._assign_usvs_to_targets()
    
    # 禁用自动重新分配
    sim.reassignment_cooldown=float('inf')
    sim.enable_reassignment=False
    
    vis = Visualizer(sim, show_metrics=True, figsize=(14, 9))
    
    if args.save_animation:
        vis.save_animation(args.save_animation, frames=400, fps=30)
    else:
        anim = vis.animate(frames=400, interval=40)
    
    print(Plotter.generate_report(sim))
    if args.save_plot:
        Plotter.plot_metrics(sim, save_path=args.save_plot)


if __name__ == "__main__":
    main() 